import networkx as nx

class SmallWorldEdgeGen(object):


    def __init__(self, numConnectedNeighbours, rewireProb):
        '''
        Constructor
        '''
        self.NumConnectedNeighbours = numConnectedNeighbours
        self.RewireProb = rewireProb
        
    def Apply(self, network):
        smallWorldGraph = nx.watts_strogatz_graph(len(network.agentList), self.NumConnectedNeighbours, self.RewireProb)

        for edge in smallWorldGraph.edges_iter():
            nodeU = network.agentList[edge[0]]
            nodeV = network.agentList[edge[1]]
            network.graph.add_edge(nodeU, nodeV)
            network.graph.add_edge(nodeV, nodeU)
